use std::{
    collections::{BTreeMap, HashMap},
    io,
};

use openmls_traits::key_store::{MlsEntityId, OpenMlsKeyStore};
use tracing::*;

use crate::errors::KeystoreError;

/// List of MLS entitiy IDs.  This should be kept in sync with the enum, because
/// it doesn't provide a way to iterate over its members.  If there are any
/// missing members in this list, it'll lead to panics.
const MLS_EIDS: &[MlsEntityId] = &[
    MlsEntityId::SignatureKeyPair,
    MlsEntityId::HpkePrivateKey,
    MlsEntityId::KeyPackage,
    MlsEntityId::PskBundle,
    MlsEntityId::EncryptionKeyPair,
    MlsEntityId::GroupState,
];

/// Gets a name we can assign to MLS entities.  This will fail to compile if new
/// members are added, reminding us to update the above table that we use for init.
fn get_entity_bucket_suffix(eid: &MlsEntityId) -> &'static str {
    match eid {
        MlsEntityId::SignatureKeyPair => "sigkeypair",
        MlsEntityId::HpkePrivateKey => "hpkeprivkey",
        MlsEntityId::KeyPackage => "keypkg",
        MlsEntityId::PskBundle => "pskbundle",
        MlsEntityId::EncryptionKeyPair => "enckp",
        MlsEntityId::GroupState => "groupstate",
    }
}

pub struct SledKeystore {
    // We use a weird `&'static str` here because `MlsEntityId` isn't `Ord`,
    // `Hash`, or even `Copy`, `Clone`, or `Eq`.
    eid_trees: BTreeMap<&'static str, sled::Tree>,
}

impl SledKeystore {
    /// Opens the keystore on an existing DB.
    pub fn open_with_db(db: &sled::Db, prefix: &str) -> sled::Result<Self> {
        // Loop through and open all the trees now so that we can easily look
        // them up when we want to use them.
        let mut eid_trees = BTreeMap::new();
        for eid in MLS_EIDS {
            let tree = open_eid_tree(db, eid, prefix)?;
            eid_trees.insert(get_entity_bucket_suffix(eid), tree);
        }

        Ok(Self { eid_trees })
    }

    fn get_entity_tree(&self, eid: MlsEntityId) -> &sled::Tree {
        let name = get_entity_bucket_suffix(&eid);
        self.eid_trees
            .get(name)
            .expect("sled keystore: missing EID namemap")
    }
}

impl OpenMlsKeyStore for SledKeystore {
    type Error = KeystoreError;

    fn store<V: openmls_traits::key_store::MlsEntity>(
        &self,
        k: &[u8],
        v: &V,
    ) -> Result<(), Self::Error>
    where
        Self: Sized,
    {
        let id = get_entity_bucket_suffix(&V::ID);
        let tree = self.get_entity_tree(V::ID);

        let mut buf = Vec::new();
        ciborium::ser::into_writer(v, &mut buf)
            .map_err(|e| KeystoreError::EncodeEntity(id, e.to_string()))?;
        tree.insert(k, buf)?;

        Ok(())
    }

    fn read<V: openmls_traits::key_store::MlsEntity>(&self, k: &[u8]) -> Option<V>
    where
        Self: Sized,
    {
        let id = get_entity_bucket_suffix(&V::ID);
        let tree = self.get_entity_tree(V::ID);

        let buf = match tree.get(k) {
            Ok(Some(buf)) => buf,
            Ok(None) => {
                warn!(%id, "queried missing entity");
                return None;
            }
            Err(e) => {
                error!(%id, err = ?e, "error fetching entity");
                return None;
            }
        };

        let buf_cur = io::Cursor::new(buf);
        let Ok(val) = ciborium::de::from_reader(buf_cur) else {
            error!(%id, "failed to parse entity");
            return None;
        };

        Some(val)
    }

    fn delete<V: openmls_traits::key_store::MlsEntity>(&self, k: &[u8]) -> Result<(), Self::Error> {
        let id = get_entity_bucket_suffix(&V::ID);
        let tree = self.get_entity_tree(V::ID);

        let v = tree.remove(k)?;
        if v.is_none() {
            warn!(%id, "deleted missing entity");
            Err(KeystoreError::MissingEntity(id))
        } else {
            Ok(())
        }
    }
}

fn open_eid_tree(db: &sled::Db, eid: &MlsEntityId, prefix: &str) -> sled::Result<sled::Tree> {
    let suffix = get_entity_bucket_suffix(eid);
    let name = format!("{prefix}_{suffix}");
    db.open_tree(name)
}
